/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * Description: Service header generator.
 */
#include "datasystem/common/rpc/plugin_generator/zmq_rpc_generator.h"

namespace datasystem {

void ZmqRpcGenerator::CreateServiceHeader(const google::protobuf::FileDescriptor &file,
                                          compiler::GeneratorContext *generatorCtx) const
{
    std::unique_ptr<io::ZeroCopyOutputStream> outputFile(generatorCtx->Open(fileName + ".service.rpc.pb.h"));
    io::Printer printer(outputFile.get(), '$');

    GenerateServicePrologue(printer, file);
    printer.PrintRaw(namespaceBegin);

    for (auto i = 0; i < file.service_count(); ++i) {
        auto *svc = file.service(i);
        GenerateServiceClass(printer, *svc, PREFIX);
    }

    printer.PrintRaw(namespaceEnd);
    printer.PrintRaw(ENDIF);
}

void ZmqRpcGenerator::GenerateServiceClass(io::Printer &printer, const google::protobuf::ServiceDescriptor &svc,
                                           const std::string &indent) const
{
    const std::string &svcName = svc.name();
    std::map<std::string, std::string> vars;
    vars["svc_name"] = svcName;

    // Go through all the methods to create the MethodServiceClass.
    GenerateMethodClass(printer, svc);

    const std::string impl =
        // Now the main body.
        "class $svc_name$_Stub;\n"
        "class $svc_name$ : public ::datasystem::ZmqService {\n"
        "public:\n"
        "    typedef $svc_name$_Stub Stub;\n"
        // ZMQ version of Constructor, common to both code paths.
        "    $svc_name$() { InitMethodMap(); serviceName_ = \"$svc_name$\"; }\n"
        "    explicit $svc_name$(datasystem::HostPort localAddress)\n"
        "    : localAddress_(std::move(localAddress)) { InitMethodMap(); serviceName_ = \"$svc_name$\"; }\n";
    printer.Print(vars, impl.c_str());

    // Go through all the methods to create the MethodServiceClass.
    GenerateMethodClass(printer, svc);

    // Create the InitMethodMap.
    GenerateInitMethodMapDecl(printer);
    // Implement the override function CallMethod.
    ImplementZmqCallMethodDecl(printer);
    // Implement the override function ServiceName.
    GenerateSvcName(printer, svcName, indent);

    std::string impl2 = "    ~$svc_name$() override = default;\n";
    printer.Print(vars, impl2.c_str());
    // One virtual function for inherited class to override.
    printer.PrintRaw("    virtual ::datasystem::Status Init() { return ::datasystem::Status::OK(); }\n");

    // State the virtual functions that user must implement.
    ListVirtualFunctions(printer, svc);

    std::string impl3 =
        "protected:\n"
        "    datasystem::HostPort localAddress_;\n"
        "    const datasystem::HostPort &GetLocalAddr() const { return localAddress_; };\n"
        "private:\n"
        "};\n";
    printer.PrintRaw(impl3);
}

void ZmqRpcGenerator::GenerateServicePrologue(io::Printer &printer, const google::protobuf::FileDescriptor &file) const
{
    std::map<std::string, std::string> vars;
    vars["full_file_name"] = file.name();
    vars["file_name"] = fileName;
    std::string impl =
        "// Generated by the protocol buffer compiler.  DO NOT EDIT!\n"
        "// source: $full_file_name$\n"
        "#ifndef DATASYSTEM_PROTO_ZMQ_$file_name$_SERVICE_H\n"
        "#define DATASYSTEM_PROTO_ZMQ_$file_name$_SERVICE_H\n"
        "#include <memory>\n";
    if (file.message_type_count() > 0) {
        impl += "#include \"$file_name$.pb.h\"\n";
    }
    impl +=
        "#include \"datasystem/common/rpc/rpc_server_stream_base.h\"\n"
        "#include \"datasystem/common/rpc/rpc_message.h\"\n"
        "#include \"datasystem/common/rpc/rpc_server.h\"\n"
        "#include \"datasystem/common/rpc/zmq/zmq_stub.h\"\n"
        "#include \"datasystem/common/rpc/zmq/zmq_service.h\"\n"
        "#include \"datasystem/common/util/net_util.h\"\n"
        "#include \"datasystem/utils/status.h\"\n";
    printer.Print(vars, impl.c_str());
    for (auto k = 0; k < file.dependency_count(); ++k) {
        auto depend = file.dependency(k)->name();
        // Check if it ends in ".proto".
        bool isAProto = StripSuffix(depend, ".proto");
        if (isAProto) {
            vars["depend"] = depend;
            printer.Print(vars, "#include \"$depend$.pb.h\"\n");
        }
    }
    printer.PrintRaw("#include <string>\n");
}

void ZmqRpcGenerator::GenerateMethodClass(io::Printer &printer, const google::protobuf::ServiceDescriptor &svc)
{
    for (auto j = 0; j < svc.method_count(); ++j) {
        if (svc.method(j) == nullptr) {
            continue;
        }
        auto &method = *(svc.method(j));
        printer.PrintRaw("class " + MethodSvcClassName(method.name())
                         + " final : public ::datasystem::RpcServiceMethod {\n");
        printer.PrintRaw("    std::string MethodName() const override { return \"" + method.name() + "\"; }\n");
        printer.PrintRaw("    int32_t MethodIndex() const override { return " + std::to_string(j) + "; }\n");
        if (HasPayloadSendOption(method)) {
            printer.PrintRaw("    bool HasPayloadSendOption() const override { return true; }\n");
        }
        if (HasPayloadRecvOption(method)) {
            printer.PrintRaw("    bool HasPayloadRecvOption() const override { return true; }\n");
        }
        if (method.client_streaming()) {
            printer.PrintRaw("    virtual bool ClientStreaming() const override { return true; }\n");
        }
        if (method.server_streaming()) {
            printer.PrintRaw("    virtual bool ServerStreaming() const override { return true; }\n");
        }
        if (UnarySocketNeeded(method)) {
            printer.PrintRaw("    bool HasUnarySocketOption() const override { return true; }\n");
        }
        printer.PrintRaw("};\n");
    }
}

void ZmqRpcGenerator::ListVirtualFunctions(io::Printer &printer, const google::protobuf::ServiceDescriptor &svc)
{
    for (auto j = 0; j < svc.method_count(); ++j) {
        if (svc.method(j) == nullptr) {
            continue;
        }
        auto &method = *(svc.method(j));
        std::string virtual_func = "    virtual ::datasystem::Status ";
        virtual_func += method.name();
        virtual_func += "(";
        if (method.client_streaming() && method.server_streaming()) {
            virtual_func += "std::shared_ptr<::datasystem::ServerWriterReader<" + method.output_type()->name() + ",";
            virtual_func += method.input_type()->name() + ">>";
        } else if (!method.client_streaming() && method.server_streaming()) {
            virtual_func += "std::shared_ptr<::datasystem::ServerWriter<" + method.output_type()->name();
            virtual_func += ">>, ";
            virtual_func += method.input_type()->name() + " &";
        } else if (method.client_streaming() && !method.server_streaming()) {
            virtual_func += "std::shared_ptr<::datasystem::ServerReader<" + method.input_type()->name();
            virtual_func += ">>, ";
            virtual_func += method.output_type()->name() + " &";
        } else if (UnarySocketNeeded(method)) {
            virtual_func +=
                "std::shared_ptr<::datasystem::ServerUnaryWriterReader<" + method.output_type()->name() + ",";
            virtual_func += method.input_type()->name() + ">>";
        } else {
            virtual_func += "const " + method.input_type()->name() + " &, ";
            virtual_func += method.output_type()->name() + " &";
        }
        if (HasPayloadSendOption(method) && !method.client_streaming() && !UnarySocketNeeded(method)) {
            virtual_func += ", std::vector<::datasystem::RpcMessage> payload";
        }
        if (HasPayloadRecvOption(method) && !method.server_streaming() && !UnarySocketNeeded(method)) {
            virtual_func += ", std::vector<::datasystem::RpcMessage> &outPayload";
        }
        virtual_func += ") = 0;\n";
        printer.PrintRaw(virtual_func);
    }
}

void ZmqRpcGenerator::ImplementZmqCallMethodDecl(io::Printer &printer)
{
    printer.PrintRaw(
        "    ::datasystem::Status CallMethod(std::shared_ptr<::datasystem::ZmqServerMsgQueRef> sock, "
        "::datasystem::MetaPb meta, std::deque<::datasystem::ZmqMessage> &&inMsg, int64_t seqNo) override;\n");
}
}  // namespace datasystem
